import cv2
import numpy as np
from datetime import datetime
import math
from scipy.interpolate import CubicSpline  # type: ignore # 用于样条曲线拟合（需安装scipy：pip install scipy）

# -------------------------- 1. 配置参数（根据蝴蝶图像调整）--------------------------
IMAGE_PATH = "butterfly_with_ruler.jpeg"  # 蝴蝶图像路径
BLUR_KERNEL = (5, 5)                  # 高斯模糊核（去噪）
CANNY_LOW = 50                        # Canny边缘检测低阈值
CANNY_HIGH = 150                      # Canny边缘检测高阈值
PIXEL_TO_MM = 0.08                     # 像素→毫米比例（需用参照物校准，例：12像素=1mm则设0.08）
EXPORT_TXT_PATH = "butterfly_contour_result.txt"  # 结果导出路径
FIT_ERROR_THRESH = 2.0                # 拟合误差阈值：>2为曲线，≤2为直线（单位：像素）
SAMPLE_INTERVAL = 2                   # 曲线段采样间隔（越小轮廓越细腻）

# -------------------------- 2. 全局变量（轮廓数据与状态）--------------------------
contour_points = []  # 手动选点的轮廓点（(x1,y1), (x2,y2), ...）
is_select_complete = False  # 是否完成轮廓选点
fitted_segments = []  # 分段拟合结果：每个元素=(类型, 起点, 终点, 长度_px, 长度_mm, 弧度, 采样点)
total_length_px = 0.0  # 轮廓总长度（像素）
total_length_mm = 0.0  # 轮廓总长度（毫米）


# -------------------------- 3. 辅助函数（预处理、拟合、计算）--------------------------
def preprocess_image(img_path):
    """图像预处理：去噪→边缘检测→增强轮廓，便于选点"""
    img = cv2.imread(img_path)
    if img is None:
        raise ValueError(f"无法读取图像：{img_path}")
    # 灰度化→高斯模糊→Canny边缘检测
    img_gray = cv2.cvtColor(img, cv2.COLOR_BGR2GRAY)
    img_blur = cv2.GaussianBlur(img_gray, BLUR_KERNEL, 0)
    img_edge = cv2.Canny(img_blur, CANNY_LOW, CANNY_HIGH)
    # 形态学膨胀：增强边缘连续性
    kernel = np.ones((2, 2), np.uint8)
    img_edge = cv2.dilate(img_edge, kernel, iterations=1)
    # 边缘叠加到原图（半透明），辅助选点
    img_with_edge = img.copy()
    img_with_edge[img_edge > 0] = [0, 255, 0]  # 边缘设为绿色
    img_with_edge = cv2.addWeighted(img, 0.7, img_with_edge, 0.3, 0)  # 半透明叠加
    return img, img_with_edge


def calculate_segment_type(start_pt, end_pt, mid_pts):
    """判断线段类型：直线/曲线（通过拟合直线的误差）"""
    # 拟合直线：用起点和终点生成直线方程 ax + by + c = 0
    x1, y1 = start_pt
    x2, y2 = end_pt
    a = y2 - y1
    b = x1 - x2
    c = x2*y1 - x1*y2
    # 计算所有中间点到直线的平均距离（误差）
    total_error = 0.0
    for (x, y) in mid_pts:
        distance = abs(a*x + b*y + c) / math.hypot(a, b)  # 点到直线的距离公式
        total_error += distance
    avg_error = total_error / len(mid_pts) if mid_pts else 0.0
    # 判断类型：误差>阈值为曲线，否则为直线
    return "曲线" if avg_error > FIT_ERROR_THRESH else "直线", avg_error


def fit_curve(points):
    """拟合曲线（三次样条），并采样轮廓点"""
    points = np.array(points)
    x = points[:, 0]
    y = points[:, 1]
    # 生成参数t（按点的顺序从0到1）
    t = np.linspace(0, 1, len(points))
    # 三次样条拟合
    cs_x = CubicSpline(t, x)
    cs_y = CubicSpline(t, y)
    # 采样曲线点（间隔SAMPLE_INTERVAL，确保轮廓细腻）
    t_sample = np.arange(0, 1 + 1/SAMPLE_INTERVAL, 1/SAMPLE_INTERVAL)
    x_sample = cs_x(t_sample)
    y_sample = cs_y(t_sample)
    sampled_points = [(round(x, 2), round(y, 2)) for x, y in zip(x_sample, y_sample)]
    # 计算曲线长度（累加采样点间的距离）
    curve_length_px = 0.0
    for i in range(1, len(sampled_points)):
        x1, y1 = sampled_points[i-1]
        x2, y2 = sampled_points[i]
        curve_length_px += math.hypot(x2 - x1, y2 - y1)
    return sampled_points, curve_length_px


def calculate_curve_radian(sampled_points):
    """计算曲线段的平均弧度（通过拟合圆的曲率）"""
    if len(sampled_points) < 3:
        return 0.0  # 点太少无法计算弧度
    # 取曲线中间3个点拟合圆，近似曲率
    mid_idx = len(sampled_points) // 2
    pts = np.array(sampled_points[mid_idx-1:mid_idx+2], dtype=np.float32)
    # 拟合圆：(x0,y0)为圆心，r为半径
    (x0, y0), r = cv2.minEnclosingCircle(pts)
    if r < 1e-3:
        return 0.0  # 半径过小，视为直线
    # 弧度=弧长/半径（取曲线段1/3长度近似弧长）
    curve_segment_len = 0.0
    for i in range(mid_idx-5, mid_idx+5):
        if 0 <= i < len(sampled_points)-1:
            x1, y1 = sampled_points[i]
            x2, y2 = sampled_points[i+1]
            curve_segment_len += math.hypot(x2 - x1, y2 - y1)
    radian = curve_segment_len / r if r != 0 else 0.0
    return round(radian, 4)


def split_contour_segments(points):
    """将完整轮廓拆分为多段，每段含2-5个点（平衡拟合精度与分段数量）"""
    segments = []
    start_idx = 0
    # 每段取2-5个点（根据轮廓复杂度自动调整）
    while start_idx < len(points) - 1:
        # 若剩余点少，取全部；否则取2-5个点
        end_idx = min(start_idx + np.random.randint(2, 6), len(points))
        segment_points = points[start_idx:end_idx]
        if len(segment_points) >= 2:
            segments.append(segment_points)
        start_idx = end_idx - 1  # 相邻段重叠1个点，避免断裂
    return segments


# -------------------------- 4. 鼠标交互函数（选点+轮廓绘制）--------------------------
def mouse_callback(event, x, y, flags, param):
    global contour_points, is_select_complete, fitted_segments, total_length_px, total_length_mm
    img_display = param["img_display"]
    img_original = param["img_original"]

    # 单击左键：添加轮廓点（未完成选点时）
    if not is_select_complete and event == cv2.EVENT_LBUTTONDOWN:
        contour_points.append((x, y))
        # 绘制选点（红色实心圆+序号）
        cv2.circle(img_display, (x, y), 3, (0, 0, 255), -1)
        cv2.putText(img_display, str(len(contour_points)), (x+5, y-5),
                    cv2.FONT_HERSHEY_SIMPLEX, 0.5, (255, 0, 0), 1)
        # 实时绘制已选轮廓（蓝色虚线）
        if len(contour_points) >= 2:
            cv2.polylines(img_display, [np.array(contour_points)], isClosed=False,
                          color=(255, 0, 0), thickness=1, lineType=cv2.LINE_AA)
        print(f"添加轮廓点{len(contour_points)}: ({x}, {y})")

    # 双击左键：结束轮廓选点，开始分段拟合与测量
    elif event == cv2.EVENT_LBUTTONDBLCLK and not is_select_complete:
        if len(contour_points) < 3:
            print("提示：至少需要3个点才能构成轮廓，请继续选点！")
            return
        is_select_complete = True
        print("\n轮廓选点完成，开始分段拟合与测量...")

        # 1. 拆分轮廓为多段（适配复杂形态）
        contour_segments = split_contour_segments(contour_points)
        fitted_segments = []
        total_length_px = 0.0

        # 2. 逐段拟合与测量
        for seg_idx, seg_points in enumerate(contour_segments, 1):
            if len(seg_points) < 2:
                continue
            start_pt = seg_points[0]
            end_pt = seg_points[-1]
            mid_pts = seg_points[1:-1] if len(seg_points) > 2 else []

            # 3. 判断线段类型（直线/曲线）
            seg_type, fit_error = calculate_segment_type(start_pt, end_pt, mid_pts)
            seg_length_px = 0.0
            seg_sampled_points = []
            seg_radian = 0.0

            # 4. 直线段处理
            if seg_type == "直线":
                # 计算直线长度
                seg_length_px = math.hypot(end_pt[0]-start_pt[0], end_pt[1]-start_pt[1])
                # 直线段采样点（起点+终点）
                seg_sampled_points = [start_pt, end_pt]
                # 直线弧度为0
                seg_radian = 0.0
                # 绘制直线段（绿色实线）
                cv2.line(img_display, start_pt, end_pt, (0, 255, 0), 2, cv2.LINE_AA)

            # 5. 曲线段处理
            else:
                # 样条曲线拟合与采样
                seg_sampled_points, seg_length_px = fit_curve(seg_points)
                # 计算曲线弧度
                seg_radian = calculate_curve_radian(seg_sampled_points)
                # 绘制曲线段（橙色实线）
                curve_pts = np.array(seg_sampled_points, dtype=np.int32)
                cv2.polylines(img_display, [curve_pts], isClosed=False,
                              color=(0, 165, 255), thickness=2, lineType=cv2.LINE_AA)

            # 6. 累加总长度，记录分段结果
            total_length_px += seg_length_px
            seg_length_mm = seg_length_px * PIXEL_TO_MM
            fitted_segments.append({
                "段号": seg_idx,
                "类型": seg_type,
                "起点": start_pt,
                "终点": end_pt,
                "长度_像素": round(seg_length_px, 2),
                "长度_毫米": round(seg_length_mm, 2),
                "弧度": seg_radian,
                "采样点数量": len(seg_sampled_points),
                "采样点": seg_sampled_points
            })

        # 7. 计算总长度（毫米）
        total_length_mm = total_length_px * PIXEL_TO_MM

        # 8. 显示测量结果（图像左上角，黑色背景防遮挡）
        result_text = [
            f"轮廓总长度：{total_length_px:.1f}px ({total_length_mm:.2f}mm)",
            f"分段数量：{len(fitted_segments)}段（绿=直线，橙=曲线）",
            f"按'S'导出结果 | 按'R'重新选点 | 按'ESC'退出"
        ]
        for i, text in enumerate(result_text):
            cv2.rectangle(img_display, (10, 20+i*30), (420, 45+i*30), (0, 0, 0), -1)
            cv2.putText(img_display, text, (15, 40+i*30),
                        cv2.FONT_HERSHEY_SIMPLEX, 0.7, (255, 255, 255), 1)

        print(f"\n测量完成！")
        print(f"轮廓总长度：{total_length_px:.1f}像素 = {total_length_mm:.2f}毫米")
        print(f"共分为{len(fitted_segments)}段（直线段+曲线段）")

    # 按R键：重新选点（重置状态）
    elif event == cv2.EVENT_KEYDOWN and event == ord('r'):
        contour_points = []
        is_select_complete = False
        fitted_segments = []
        total_length_px = 0.0
        total_length_mm = 0.0
        # 重置显示图像（恢复边缘叠加的原图）
        img_display[:] = preprocess_image(IMAGE_PATH)[1]
        print("\n已重置，可重新选点！")


# -------------------------- 5. 结果导出函数（完整轮廓数据）--------------------------
def export_contour_result():
    if not is_select_complete or len(fitted_segments) == 0:
        print("提示：未完成轮廓测量，无法导出结果！")
        return
    with open(EXPORT_TXT_PATH, "w", encoding="utf-8") as f:
        f.write(f"蝴蝶轮廓测量报告 - {datetime.now().strftime('%Y-%m-%d %H:%M:%S')}\n")
        f.write(f"图像路径：{IMAGE_PATH}\n")
        f.write(f"像素→毫米比例：1px = {PIXEL_TO_MM}mm\n")
        f.write(f"轮廓总长度：{total_length_px:.2f}px = {total_length_mm:.2f}mm\n")
        f.write(f"分段数量：{len(fitted_segments)}段\n")
        f.write("\n" + "="*80 + "\n")

        # 1. 导出手动选点的完整轮廓
        f.write("一、手动选点轮廓（序号：x,y）\n")
        for i, (x, y) in enumerate(contour_points, 1):
            f.write(f"{i}: {x}, {y}\n")
        f.write("\n" + "="*80 + "\n")

        # 2. 导出分段测量结果
        f.write("二、分段测量详情\n")
        for seg in fitted_segments:
            f.write(f"\n段号：{seg['段号']} | 类型：{seg['类型']}\n")
            f.write(f"起点：{seg['起点']} | 终点：{seg['终点']}\n")
            f.write(f"长度：{seg['长度_像素']}px = {seg['长度_毫米']}mm | 弧度：{seg['弧度']}\n")
            f.write(f"采样点数量：{seg['采样点数量']} | 采样点：{seg['采样点']}\n")
        f.write("\n" + "="*80 + "\n")

        # 3. 导出完整轮廓的采样点（可直接用于建模）
        f.write("三、完整轮廓采样点（按顺序，共{}个点）\n".format(
            sum(seg["采样点数量"] for seg in fitted_segments)
        ))
        point_idx = 1
        for seg in fitted_segments:
            for (x, y) in seg["采样点"]:
                f.write(f"{point_idx}: {x}, {y} | 所属段：{seg['段号']}（{seg['类型']}）\n")
                point_idx += 1
    print(f"\n结果已导出至：{EXPORT_TXT_PATH}")


# -------------------------- 6. 主程序（入口）--------------------------
if __name__ == "__main__":
    try:
        # 1. 图像预处理（增强轮廓，辅助选点）
        img_original, img_display = preprocess_image(IMAGE_PATH)

        # 2. 初始化显示窗口与鼠标回调
        cv2.namedWindow("Butterfly Contour Measurement", cv2.WINDOW_NORMAL)
        cv2.resizeWindow("Butterfly Contour Measurement", 1200, 800)  # 调整窗口大小
        # 传递原图与显示图到回调函数（避免全局变量污染）
        callback_param = {"img_original": img_original, "img_display": img_display}
        cv2.setMouseCallback("Butterfly Contour Measurement", mouse_callback, callback_param)

        # 3. 显示初始提示
        init_text = [
            "蝴蝶轮廓测量程序",
            "操作说明：",
            "1. 单击左键：添加轮廓点（沿蝴蝶边缘选点）",
            "2. 双击左键：结束选点并开始测量",
            "3. 按'S'：导出测量结果",
            "4. 按'R'：重新选点",
            "5. 按'ESC'：退出程序"
        ]
        for i, text in enumerate(init_text):
            y_pos = 50 + i*30
            cv2.putText(img_display, text, (10, y_pos),
                        cv2.FONT_HERSHEY_SIMPLEX, 0.7, (0, 0, 255), 1)

        # 4. 主循环（等待按键操作）
        print("程序启动成功！按提示沿蝴蝶边缘选点...")
        while True:
            cv2.imshow("Butterfly Contour Measurement", img_display)
            key = cv2.waitKey(1) & 0xFF
            if key == 27:  # ESC键：退出
                print("\n程序退出！")
                break
            elif key == ord('s'):  # S键：导出结果
                export_contour_result()
            elif key == ord('r'):  # R键：重新选点
                contour_points = []
                is_select_complete = False
                fitted_segments = []
                total_length_px = 0.0
                total_length_mm = 0.0
                img_display[:] = preprocess_image(IMAGE_PATH)[1]  # 重置显示图
                print("\n已重置，可重新选点！")

    except Exception as e:
        print(f"程序异常：{str(e)}")
    finally:
        cv2.destroyAllWindows()